import qutip as qt
import numpy as np
"""
Suppose A as a⁺
Suppose the first qubit is at the right side of ket |3,2,1>
"""
π = np.pi
GHz = 1e9
MHz = 1e6
ns = 1e-9


def to_base(dig, base):
    BS = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    res = ''
    while dig:
        res += BS[dig % base]
        dig //= base
    res = res.ljust(base, '0')
    return res[::-1] or '0' * base


def aᵢ(i, l=3, n_qubits=3):
    """
                aᵢ
    aᵢ for l level n qubits system.
    """
    ops = [qt.identity(l)] * n_qubits
    ops[-i] = qt.destroy(l)
    return qt.tensor(ops)


def Aᵢ(i, l=3, n_qubits=3):
    """
                aᵢ⁺
    aᵢ⁺ for l level n qubits system.
    """
    ops = [qt.identity(l)] * n_qubits
    ops[-i] = qt.create(l)
    return qt.tensor(ops)


def Aᵢaⱼ(i, j, l=3, n_qubits=3):
    """
                aᵢ⁺aⱼ
    """
    return Aᵢ(i, l, n_qubits) * aᵢ(j, l, n_qubits)


def aᵢAⱼ(i, j, l=3, n_qubits=3):
    """
                aᵢaⱼ⁺
    """
    return aᵢ(i, l, n_qubits) * Aᵢ(j, l, n_qubits)


def AᵢAᵢaᵢaᵢ(i, l=3, n_qubits=3):
    """
                aᵢ⁺aᵢ⁺aᵢaᵢ
    """
    return Aᵢ(i, l, n_qubits) * Aᵢaⱼ(i, i, l, n_qubits) * aᵢ(i, l, n_qubits)


def Aᵢaⱼ_plus_aᵢAⱼ(i, j, l=3, n_qubits=3):
    """
                aᵢ⁺aⱼ + aᵢaⱼ⁺
    """
    return Aᵢaⱼ(i, j, l, n_qubits) + aᵢAⱼ(i, j, l, n_qubits)


def Ham(ω1, ω2, ωc, α1, α2, αc, g12, g1c, g2c):
    first = ω1 * Aᵢaⱼ(1, 1) + ω2 * Aᵢaⱼ(2, 2) + ωc * Aᵢaⱼ(3, 3)
    second = α1 * AᵢAᵢaᵢaᵢ(1) + α2 * AᵢAᵢaᵢaᵢ(2) + αc * AᵢAᵢaᵢaᵢ(3)
    second /= 2
    third = g12 * Aᵢaⱼ_plus_aᵢAⱼ(1, 2) + g1c * Aᵢaⱼ_plus_aᵢAⱼ(
        1, 3) + g2c * Aᵢaⱼ_plus_aᵢAⱼ(2, 3)
    return first + second + third


def sin_pulse(t, T, ωc_max, ωc_min):
    ωc_freq = 2 * π / T
    return (ωc_max - ωc_min) * (np.sin(ωc_freq * t) + 1) / 2 + ωc_min


def Uz_propagator(ham, pulse, time_range):
    dt = time_range[1] - time_range[0]
    U = 1
    for t in time_range:
        U = (-1j * ham(pulse(t)) * dt).expm() * U
    return U


def vector_spectrum(num_energy_levels, hamiltonian, n_qubits=3):
    vaule, vector = hamiltonian.eigenstates()
    index_energy = {}
    for i in range(num_energy_levels):
        index = np.argmax(np.abs(np.array(vector[i])))
        bit_string = to_base(index, n_qubits)
        if bit_string in index_energy:
            index_energy[bit_string].append(vector[i])
        else:
            index_energy[bit_string] = [vector[i]]
    return index_energy


def fidelity_metric(phase, U, res, n_qubits, U_ideal):
    U_qubit = np.matrix(np.zeros((4, 4)).astype(np.complex))
    for i in range(4):
        for j in range(4):
            basis_i = bin(i)[2:].zfill(n_qubits)
            basis_j = bin(j)[2:].zfill(n_qubits)
            U_qubit[i, j] = (res[basis_i][-1].dag() * U *
                             res[basis_j][-1]).data.toarray()[0, 0]
    U_qubit = U_qubit * (
        -1j * phase[0] * qt.tensor(qt.identity(2), -qt.sigmaz())).expm() * (
            -1j * phase[1] * qt.tensor(-qt.sigmaz(), qt.identity(2))).expm()
    F = np.real((np.trace(U_qubit.H * U_qubit) +
                 (np.abs(np.trace(U_ideal.dag() * U_qubit))**2)) / 20)
    return 1 - F


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from scipy.optimize import minimize

    ω1 = 2 * π * 5.27 * GHz
    ω2 = 2 * π * 4.62 * GHz
    ω3 = ωc = 2 * π * 6.74 * GHz
    ωc_max = 2 * π * 6.74 * GHz
    ωc_min = 2 * π * 5.2 * GHz
    α1 = 2 * π * (-210) * MHz
    α2 = 2 * π * (-240) * MHz
    α3 = αc = 2 * π * (-370) * MHz
    g12 = 2 * π * 12 * MHz
    g1c = 2 * π * 122 * MHz
    g2c = 2 * π * 105 * MHz
    b0 = qt.basis(3, 0)
    b1 = qt.basis(3, 1)
    b2 = qt.basis(3, 2)
    z3 = b0 * b0.dag() - b1 * b1.dag() + b2 * b2.dag()
    cz3 = qt.tensor(
        (b0 * b0.dag() + b2 * b2.dag()), qt.identity(3)) + qt.tensor(
            b1 * b1.dag(), z3)

    u_ideal = qt.qip.operations.controlled_gate(qt.sigmaz())

    # def fun(params):
    T = 60 * ns
    # phase = params[:]
    ham = lambda ωc: Ham(ω1, ω2, ωc, α1, α2, αc, g12, g1c, g2c)
    pulse = lambda t: sin_pulse(t, T, ωc_min, ωc_max)
    time_range = np.arange(0, T / 2, 0.1 * ns)
    uz = Uz_propagator(ham, pulse, time_range)
    vs = vector_spectrum(10, ham(ωc_max))
    # infid = fidelity_metric(phase, uz, vs, 3, u_ideal)
    # print(infid)
    # return infid

    bounds = ((-2 * π, 2 * π), (-2 * π, 2 * π))
    res = minimize(fidelity_metric, (0.15, 1.5),
                   args=(uz, vs, 3, u_ideal),
                   method='SLSQP',
                   bounds=bounds,
                   tol=1e-7,
                   options={'disp': True})
    fid = 1 - res.fun
    print('fid: {:.3f}%'.format(fid * 100))
